import numpy as np
from sklearn.manifold import TSNE
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import time
import json
import pdb

meta_type = 1

t0 = time.time()
file1 = 'model_path/pooled_feat/tsne_data.json'
data1 = json.load(open(file1, 'r'))
tsne = np.array(data1['tsne_data'])
cls_gt = np.array(data1['cls_gt'])
cls_vector = np.array(data1['cls_vector'])
t1 = time.time()
print('After %.2f seconds, loaded the whole data.' % (t1 - t0))

color = ['#66FFFF',
'#808080',
'#00FFFF',
'#7FFFD4',
'#008000',
'#90EE90',
'#DC143C',
'#FFA500',
'#FFEBCD',
'#0000FF',
'#8A2BE2',
'#A52A2A',
'#DEB887',
'#5F9EA0',
'#FF7F50',
'#D2691E',
'#7FFF00',
'#6495ED',
'#000000',
'#D3D3D3',
'#FFD700']

xmin = tsne[:, 0].min()
xmax = tsne[:, 0].max()
ymin = tsne[:, 1].min()
ymax = tsne[:, 1].max()


for score in range(10, 0, -1):
    for cls in range(1, 21):
    
        keep = cls_gt[:, 0] == cls
        x = tsne[keep, 0]
        y = tsne[keep, 1]
        cls_score = cls_vector[keep, :]

        keep = (cls_score[:, cls] >= (score-1)/10) & (cls_score[:, cls] < score/10)
        
        x_score = x[keep]
        y_score = y[keep]
        if cls < 16:
            plt.scatter(x_score, y_score, s=1, c = color[cls], alpha = 1)
        else:
            plt.scatter(x_score, y_score, s=1, marker='x', c = color[cls], alpha = 1)

t2 = time.time()
print('scatter plot time used: %.2f' % (t2 - t1))
if meta_type == 1:
    plt.legend(['aeroplane', 'bicycle', 'boat', 'bottle', 'car', 'cat', 'chair', 'diningtable', 'dog', 'horse', 'person', 'pottedplant', 'sheep', 'train', 'tvmonitor', 'bird', 'bus', 'cow', 'motorbike', 'sofa'], bbox_to_anchor=(1, 1))
elif meta_type == 2:
    plt.legend(['bicycle', 'bird', 'boat', 'bus', 'car', 'cat', 'chair', 'diningtable', 'dog', 'motorbike',
                         'person', 'pottedplant', 'sheep', 'train', 'tvmonitor', 'aeroplane', 'bottle', 'cow',
                         'horse', 'sofa'], bbox_to_anchor=(1, 1))
elif meta_type == 3:
    plt.legend(['aeroplane', 'bicycle', 'bird', 'bottle', 'bus', 'car', 'chair', 'cow', 'diningtable',
                         'dog', 'horse', 'person', 'pottedplant', 'train', 'tvmonitor', 'boat', 'cat', 'motorbike',
                         'sheep', 'sofa'], bbox_to_anchor=(1, 1))

plt.xlim(xmin, xmax)
plt.ylim(ymin, ymax)
plt.tight_layout()
plt.savefig('tsne.jpg', dpi=1000)
plt.cla()